使用Blade优化输入为Dynamic Shape的模型_人工智能平台 PAI(PAI)-阿里云帮助中心

常规推理优化普遍针对输入为Static Shape的模型，如果实际推理的模型Shape发生变化，推理优化效果就可能失效。在实际生产中，输入为Dynamic Shape的模型越来越多，因此对不同输入Shape的推理过程具有强烈的优化需求。本文介绍如何使用Blade优化输入为Dynamic Shape的模型。

使用限制

本文使用的环境需要满足以下版本要求：

系统环境：Linux系统中使用Python 3.6及其以上版本。
框架：PyTorch 1.7.1。
设备及后端：NVIDIA T4、CUDA 11.0。
推理优化工具：Blade 3.17.0及其以上版本。

操作流程

使用Blade优化输入为Dynamic Shape的ResNet50流程如下：

步骤一：准备工作
构建测试数据和模型，本文使用torchvision中标准的ResNet50模型。
步骤二：配置用于优化的config
根据Dynamic Shape的范围配置Blade config。
步骤三：调用Blade优化模型
调用blade.optimize接口优化模型，并保存优化后的模型。
步骤四：验证性能与正确性
对优化前后的推理速度及推理结果进行测试，从而验证优化报告中信息的正确性。
步骤五：加载运行优化后的模型
集成Blade SDK，加载优化后的模型进行推理。

步骤一：准备工作

下载模型预训练参数与测试数据。

预训练参数选自torchvision，为了加速下载过程，已将其存储至OSS中。测试数据随机选自ImageNet-1k验证集，预处理操作已完成，您可以下载后直接使用。

wget http://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/share/dynamic_ranges_pratice/resnet50-19c8e357.pth -O resnet50-19c8e357.pth
wget http://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/share/dynamic_ranges_pratice/imagenet_val_example.pt -O imagenet_val_example.pt

定义模型、加载模型参数和测试数据，并生成TorchScript。

import torch
import torchvision

# 构建Resnet50。
model = torchvision.models.resnet50().eval().cuda()
# 加载预训练参数。
ckpt = torch.load('resnet50-19c8e357.pth')
model.load_state_dict(ckpt)
# 加载测试数据。
example_input = torch.load('imagenet_val_example.pt').cuda()
# 生成TorchScript。
traced_model = torch.jit.trace(model, example_input).cuda().eval()

步骤二：配置用于优化的config

根据Dynamic Shape的范围配置Blade config，Blade支持任意维度的动态范围。本文以Batch维度演示config的配置。

定义Dynamic Shape的范围。
一组有效的动态范围，需要包括以下三个字段：
- min：表示dynamic shape的下界。
- max：表示dynamic shape的上界。
- opts：表示需要特别优化的Shape，可以设置多个。通常优化后的模型在这些Shape上的推理加速比更高。
上述三个字段需要符合以下规则：
- min、max及opts中的每组Shape的长度相等，且等于网络的输入数量。
- min、max及opts中的每组Shape对应位置的数值需要满足min_num <= opt_num <= max_num。
例如构建如下Dynamic Shape的范围。
```
shapes = {
    "min": [[1, 3, 224, 224]],
    "max": [[10, 3, 224, 224]],
    "opts": [
        [[5, 3, 224, 224]],
        [[8, 3, 224, 224]],
    ]
}
```
此外，Blade支持设置多个动态范围。如果Dynamic Shape的上界和下界范围过大，可能会导致优化后的模型加速不明显，您可以将一个大的范围拆分为多个小范围，通常能够带来更好的加速效果。关于如何设置多个动态范围，请参见下文的附录：设置多个动态范围。

通过定义好的Dynamic Shape范围构建Blade config。

import blade
import blade.torch as blade_torch

# Blade Torch相关config，用于设置Dynamic Shapes。
blade_torch_cfg = blade_torch.Config()
blade_torch_cfg.dynamic_tuning_shapes = shapes

# Blade相关config，用于关闭FP16的精度检查，以获得最好的加速效果。
gpu_config = {
    "disable_fp16_accuracy_check": True,
}
blade_config = blade.Config(
    gpu_config=gpu_config
)

步骤三：调用Blade优化模型

调用blade.optimize对模型进行优化，示例代码如下。关于该接口的详细描述，请参见Python接口文档。
```
with blade_torch_cfg:
    optimized_model, _, report = blade.optimize(
        traced_model,          # 模型路径。
        'o1',                  # o1无损优化。
        config=blade_config,
        device_type='gpu',     # 面向GPU设备优化，
        test_data=[(example_input,)]  # 测试数据。
    )
```
优化模型时，您需要注意以下事宜：
- blade.optimize的第一个返回值为优化后的模型，其数据类型与输入的模型相同。在这个示例中，输入的是TorchScript，返回的是优化后的TorchScript。
- 您需要确保输入的test_data在定义的Dynamic Shape范围内。

优化完成后，打印优化报告。

print("Report: {}".format(report))

打印的优化报告类似如下输出。

Report: {
  "software_context": [
    {
      "software": "pytorch",
      "version": "1.7.1+cu110"
    },
    {
      "software": "cuda",
      "version": "11.0.0"
    }
  ],
  "hardware_context": {
    "device_type": "gpu",
    "microarchitecture": "T4"
  },
  "user_config": "",
  "diagnosis": {
    "model": "unnamed.pt",
    "test_data_source": "user provided",
    "shape_variation": "undefined",
    "message": "Unable to deduce model inputs information (data type, shape, value range, etc.)",
    "test_data_info": "0 shape: (1, 3, 224, 224) data type: float32"
  },
  "optimizations": [
    {
      "name": "PtTrtPassFp16",
      "status": "effective",
      "speedup": "4.06",
      "pre_run": "6.55 ms",
      "post_run": "1.61 ms"
    }
  ],
  "overall": {
    "baseline": "6.54 ms",
    "optimized": "1.61 ms",
    "speedup": "4.06"
  },
  "model_info": {
    "input_format": "torch_script"
  },
  "compatibility_list": [
    {
      "device_type": "gpu",
      "microarchitecture": "T4"
    }
  ],
  "model_sdk": {}
}

从优化报告可以看出本示例的优化中，PtTrtPassFp16优化项生效，带来了约4.06倍左右的加速，将模型在测试数据上的推理耗时从6.55 ms下降到了1.61 ms。上述优化结果仅为本示例的测试结果，您的优化效果以实际为准。关于优化报告的字段详情请参见优化报告。

调用PyTorch的相关函数保存并加载优化后的TorchScript模型。

file_name = "resnet50_opt.pt"
# 将优化后的模型保存到本地。
torch.jit.save(optimized_model, file_name)
# 从硬盘中加载优化后的模型。
optimized_model = torch.jit.load(file_name)

步骤四：验证性能与正确性

优化完成后，通过Python脚本对优化报告的信息进行验证。

定义benchmark方法，对模型进行10次预热，然后运行100次，最终取平均的推理时间作为推理速度。

import time

@torch.no_grad()
def benchmark(model, test_data):
    # 切换模型至验证模式。
    model = model.eval()
    
    # 预热。
    for i in range(0, 10):
        model(test_data)
        
    # 开始计时运行。
    num_runs = 100
    start = time.time()
    for i in range(0, num_runs):
        model(test_data)
    torch.cuda.synchronize()
    elapsed = time.time() - start
    rt_ms = elapsed / num_runs * 1000.0
    
    # 打印结果。
    print("{:.2f} ms.".format(rt_ms))
    return rt_ms

定义一系列不同Shape的测试数据。

dummy_inputs = []
batch_num = [1, 3, 5, 7, 9]
for n in batch_num:
    dummy_inputs.append(torch.randn(n, 3, 224, 224).cuda())

遍历每组测试数据，分别调用benchmark方法对优化前与优化后的模型进行测试，并打印结果。

for inp in dummy_inputs:
    print(f'--------------test with shape {list(inp.shape)}--------------')
    print("  Origin model inference cost:     ", end='')
    origin_rt = benchmark(traced_model, inp)
    print("  Optimized model inference cost:  ", end='')
    opt_rt = benchmark(optimized_model, inp)
    speedup = origin_rt / opt_rt
    print('  Speed up: {:.2f}'.format(speedup))
    print('')

系统返回如下类似结果。

--------------test with shape [1, 3, 224, 224]--------------
  Origin model inference cost:     6.54 ms.
  Optimized model inference cost:  1.66 ms.
  Speed up: 3.94

--------------test with shape [3, 3, 224, 224]--------------
  Origin model inference cost:     10.79 ms.
  Optimized model inference cost:  2.40 ms.
  Speed up: 4.49

--------------test with shape [5, 3, 224, 224]--------------
  Origin model inference cost:     16.27 ms.
  Optimized model inference cost:  3.25 ms.
  Speed up: 5.01

--------------test with shape [7, 3, 224, 224]--------------
  Origin model inference cost:     22.62 ms.
  Optimized model inference cost:  4.39 ms.
  Speed up: 5.16

--------------test with shape [9, 3, 224, 224]--------------
  Origin model inference cost:     28.83 ms.
  Optimized model inference cost:  5.25 ms.
  Speed up: 5.49

从结果可以看出对于不同Shape的测试数据，优化后模型的推理速度是原始模型的3.94~5.49倍。上述优化结果仅为本示例的测试结果，您的优化效果以实际为准。

使用准备工作阶段准备的真实测试数据example_input，验证优化模型的正确性。

origin_output = traced_model(example_input)
_, pred = origin_output.topk(1, 1, True, True)
print("origin model output: {}".format(pred))
opt_output = optimized_model(example_input)
_, pred = origin_output.topk(1, 1, True, True)
print("optimized model output: {}".format(pred))

系统返回如下类似结果。

origin model output: tensor([[834]], device='cuda:0')
optimized model output: tensor([[834]], device='cuda:0')

从上述结果可以看出优化前后模型对于测试数据example_input的预测均为第834类。

步骤五：加载运行优化后的模型

完成验证后，您需要对模型进行部署，Blade提供了Python和C++两种运行时SDK供您集成。关于C++的SDK使用方法请参见使用SDK部署TensorFlow模型推理，下文主要介绍如何使用Python SDK部署模型。

可选：在试用阶段，您可以设置如下的环境变量，防止因为鉴权失败而程序退出。
```
export BLADE_AUTH_USE_COUNTING=1
```
获取鉴权。
```
export BLADE_REGION=<region>
export BLADE_TOKEN=<token>
```
您需要根据实际情况替换以下参数：
- <region>：Blade支持的地域，需要加入Blade用户群获取该信息，用户群的二维码详情请参见获取Token。
- <token>：鉴权Token，需要加入Blade用户群获取该信息，用户群的二维码详情请参见获取Token。

加载运行优化后的模型。

除了增加一行import blade.runtime.torch，您无需为Blade的接入编写额外代码，即原有的推理代码无需任何改动。

import torch
import blade.runtime.torch
# <your_optimized_model_path>替换为优化后的模型路径。
opt_model_dir = <your_optimized_model_path>
# <your_infer_data>替换为用于推理的数据。
infer_data = <your_infer_data>

model = torch.jit.load(opt_model_dir)
output = model(infer_data)

附录：设置多个动态范围

如果Dynamic Shape的上界和下界范围过大，可能会导致优化后的模型加速不明显，您可以将一个大的范围拆分为多个小范围，通常能够带来更好的加速效果。例如设置如下Dynamic Shape。

shapes1 = {
    "min": [[1, 3, 224, 224]],
    "max": [[5, 3, 224, 224]],
    "opts": [
        [[5, 3, 224, 224]],
    ]
}
shapes2 = {
    "min": [[5, 3, 224, 224]],
    "max": [[10, 3, 224, 224]],
    "opts": [
        [[8, 3, 224, 224]],
    ]
}
shapes = [shapes1, shapes2]

您可以使用该shapes配置上述提及的优化config，详情请参见步骤二：配置用于优化的config。